package ch.epfl.lara.matcheck.util;

import ch.epfl.lara.matcheck.ast._



object Substitute extends Environment with ShapeKeeper {
  
  def apply(ctx: Env, e: Expr): Expr = {
    def subst(name: String): String = substitute(ctx,name)
    
    def iter[T<:Expr](e: T): T = (e match {
      case Ident(name)            => cast[T,Ident]        (Ident(new Name[VarSymbol](subst(name.getName),name.getSym)))
      case Call(obj,method,args)  => cast[T,Call]         (Call(iter(obj),method,args.map(iter(_))))
      case Select(obj,field)      => cast[T,Select]       (Select(iter(obj),field))
      case If(cond,then,elze)     => cast[T,If]           (If(iter(cond),iter(then),iter(elze)))
      case New(name,args)         => cast[T,New]          (New(name,args.map(iter(_))))
      case IsInstanceOf(expr,tpe) => cast[T,IsInstanceOf] (IsInstanceOf(iter(expr),tpe))
      case +(left,right)          => cast[T,+]            (new +(iter(left),iter(right)))
      case -(left,right)          => cast[T,-]            (new -(iter(left),iter(right)))
      case *(left,right)          => cast[T,*]            (new *(iter(left),iter(right)))
      case /(left,right)          => cast[T,/]            (new /(iter(left),iter(right)))
      case %(left,right)          => cast[T,%]            (new %(iter(left),iter(right)))
      case ==(left,right)         => cast[T,==]           (new ==(iter(left),iter(right)))
      case <(left,right)          => cast[T,<]            (new <(iter(left),iter(right)))
      case <=(left,right)         => cast[T,<=]           (new <=(iter(left),iter(right)))
      case &&(left,right)         => cast[T,&&]           (new &&(iter(left),iter(right)))
      case ||(left,right)         => cast[T,||]           (new ||(iter(left),iter(right)))
      case !(f)                   => cast[T,!]            (new !(iter(f)))
      case ->(left,right)         => cast[T,->]           (new ->(iter(left),iter(right)))
      case m : Match              => cast[T,Match]        (this(ctx,m))
      case True  | Number(_) |
           False | StringLit(_) |
           ObjectId(_)             => e 
    }).setPos(e.pos)
    iter(e)
  }
  
  
  
  
  def apply[T<:Tree with PatternMatching](ctx: Env, pm: T): T = {
    def subst(name: String): String = substitute(ctx,name)
    
    def iter[T<:Tree with PatternMatching](pm: T): T = (pm match {
      case Match(scrut,tpe,cases)          => cast[T,Match]         (Match(this(ctx,scrut),tpe,cases.map(iter(_))))
      case CaseDef(pos,pattern,guard,body) =>
        def bindings(p: PatternDef): List[String] = p match {
          case Wildcard                    => Nil
          case VarPattern(name,tpe)        => List(name.getName)
          case ClassPattern(objname,args)  => args.flatMap(bindings(_))
        }
        
        val new_bindings    = bindings(pattern)
        val new_context     = ctx -- (new_bindings) // new_bindings shade the old ones (if exists)
        val updated_guard   = guard match { 
                                case Some(g) => Some(cast[Expr,Expr] (this(new_context,g))) 
                                case None    => guard
                              }
        val updated_body    = cast[Expr,Expr]  (this(new_context,body))
        cast[T,CaseDef]   (CaseDef(pos,pattern,updated_guard,updated_body))
        
      case Wildcard                       => pm
      case ClassPattern(name,args)        => cast[T,ClassPattern]   (ClassPattern(name,args.map(iter(_))))   
      case VarPattern(name,tpe)           => cast[T,VarPattern]     (VarPattern(new Name[LocalSymbol](subst(name.getName),name.getSym),tpe))
    }).setPos(pm.pos)
    
    iter(pm)
  }
  
  
  def apply(x: String, s: String, f: Formula): Formula =
    this(empty_map + ((x,s)),f)
 
    
  def apply(ctx: List[(String,String)], f: Formula): Formula = 
    this(empty_map ++ ctx,f)
  
  def apply(ctx: Env, f: Formula): Formula = {
    def subst(name: String): String = substitute(ctx,name)
    
    def iter[T<:Formula](f: T): T = f match {
      case And(left,right)          => cast[T,And]           (And(iter(left),iter(right)))
      case Or(left,right)           => cast[T,Or]            (Or(iter(left),iter(right)))
      case Neg(f)                   => cast[T,Neg]           (Neg(iter(f)))
      case Imply(left,right)        => cast[T,Imply]         (Imply(iter(left),iter(right)))
      case Equality(left,right)     => cast[T,Equality]      (Equality(iter(left),iter(right)))
      case Forall(x,f)              => cast[T,Forall]        (Forall(x,iter(f)))
      case SubsetEq(left,right)     => cast[T,SubsetEq]      (SubsetEq(iter(left),iter(right)))
      case EqSet(left,right)        => cast[T,EqSet]         (EqSet(iter(left),iter(right)))
      case In(x,set)                => cast[T,In]            (In(iter(x),iter(set)))
      case Union(left,right)        => cast[T,Union]         (Union(iter(left),iter(right)))
      case Intersection(left,right) => cast[T,Intersection]  (Intersection(iter(left),iter(right)))
      case Eq(left,right)           => cast[T,Eq]            (Eq(iter(left),iter(right)))
      case NotEq(left,right)        => cast[T,NotEq]         (NotEq(iter(left),iter(right)))
      case Less(left,right)         => cast[T,Less]          (Less(iter(left),iter(right)))
      case LessEq(left,right)       => cast[T,LessEq]        (LessEq(iter(left),iter(right)))
      case Plus(left,right)         => cast[T,Plus]          (Plus(iter(left),iter(right)))
      case Minus(left,right)        => cast[T,Minus]         (Minus(iter(left),iter(right)))
      case Div(left,right)          => cast[T,Div]           (Div(iter(left),iter(right)))
      case Prod(left,right)         => cast[T,Prod]          (Prod(iter(left),iter(right)))
      case Mod(left,right)          => cast[T,Mod]           (Mod(iter(left),iter(right)))
      case Element(x)               => cast[T,Element]       (Element(subst(x)))
      case Deref(obj,field)         => cast[T,Deref]         (Deref(iter(obj),iter(field)))
      case Set(name)                => cast[T,Set]           (Set(subst(name)))
      case Function(name,elems)     => cast[T,Function]      (Function(subst(name),elems.map(iter(_))))
      case Tru  | EmptySet |  
           Flse | Const(_) | 
           StringLit0(_)            => f
    }
    iter(f)
  }
  
  
  
  private def substitute(ctx: Env, name: String): String = ctx.getOrElse(name,name)
  
  
  override def toString = "Substitute"
  
}